'''
https://leetcode.cn/problems/balanced-binary-tree/
'''
from typing import Optional

from c06_tree.utils import TreeNode


# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def isBalanced(self, root: Optional[TreeNode]) -> bool:
        def process(root) -> (bool, int):
            if root is None: return (True, 0)
            left_info = process(root.left)
            right_info = process(root.right)
            is_balanced = left_info[0] and right_info[0] and abs(left_info[1] - right_info[1]) <= 1
            height = 1 + max(left_info[1], right_info[1])
            return (is_balanced, height)
        return process(root)[0]

